package com.aos.iot.socket.common.handler;

import cn.hutool.core.thread.ThreadUtil;
import cn.hutool.json.JSONUtil;
import com.aos.iot.socket.common.ack.AckManager;
import com.aos.iot.socket.common.annotation.TcpEventListenerHub;
import com.aos.iot.socket.common.ack.AckCallback;
import com.aos.iot.socket.common.channel.IChannel;
import com.aos.iot.socket.common.client.ClientBoxFactor;
import com.aos.iot.socket.common.client.IClientBox;
import com.aos.iot.socket.common.message.Message;
import com.aos.iot.socket.common.protocol.IProtocol;
import com.aos.iot.socket.common.store.Store;
import com.aos.iot.socket.common.store.StoreFactory;
import com.aos.iot.socket.common.store.pubsub.PubSubStore;
import io.netty.channel.Channel;
import io.netty.channel.ChannelHandlerContext;
import io.netty.util.internal.PlatformDependent;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;

import java.lang.reflect.ParameterizedType;
import java.net.SocketAddress;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;

@Data
@Slf4j
public abstract class SocketClient<M extends Message> implements IChannel<M>, Store {

    private final ChannelHandlerContext channelHandlerContext;

    private final IProtocol<M> protocol;

    private final PubSubStore<M> pubSubStore;

    private final Store store;

    private final AckManager ackManager;

    private final ClientBoxFactor clientBoxFactor;

    private final IClientBox<M, SocketClient<M>> clientBox;

    private final Class<M> messageClass = (Class<M>)((ParameterizedType) this.getClass().getGenericSuperclass()).getActualTypeArguments()[0];

    private final Set<String> callbackFlag = ConcurrentHashMap.newKeySet();

    private final Map<String, Message> callbackReturn = PlatformDependent.newConcurrentHashMap();

    public SocketClient(ChannelHandlerContext channelHandlerContext, IProtocol<M> protocol, StoreFactory<M> storeFactory, AckManager ackManager, ClientBoxFactor clientBoxFactor) {
        this.channelHandlerContext = channelHandlerContext;
        this.protocol = protocol;
        this.pubSubStore = storeFactory.pubSubStore();
        this.store = storeFactory.createStore(channelHandlerContext.channel().id().asLongText());
        this.ackManager = ackManager;
        this.clientBoxFactor = clientBoxFactor;
        this.clientBox = clientBoxFactor.getClientBox(protocol.name());
    }

    public void disconnect() {
        close();
    }

    @Override
    public String getChannelId() {
        return channelHandlerContext.channel().id().asLongText();
    }

    @Override
    public Channel getChannel() {
        return channelHandlerContext.channel();
    }

    @Override
    public IProtocol<M> getProtocol() {
        return protocol;
    }

    @Override
    public SocketAddress getRemoteAddress() {
        return clientBox.getRemoteAddress(channelHandlerContext.channel());
    }

    @Override
    public boolean isChannelOpen() {
        return channelHandlerContext.channel().isOpen();
    }

    @Override
    public void close() {
        channelHandlerContext.channel().close();
    }

    @Override
    public void joinCluster(String clusterName) {
        pubSubStore.subscribe(clusterName, message -> {
            log.info("Topic [{}] subscribe : {}", clusterName, JSONUtil.toJsonStr(message));
            TcpEventListenerHub.invoke(protocol.name(), message.getCmd(),
            channelHandlerContext, channelHandlerContext.channel(),
            channelHandlerContext.channel().id(), protocol, message, this, message.getDirection());
        }, messageClass);
    }

    @Override
    public void leaveCluster(String clusterName) {
        pubSubStore.unsubscribe(clusterName);
    }

    @Override
    public void send(M message, AckCallback<? extends Message> ackCallback) {
        message.setDirection(Message.Direction.DOWN);
        if(ackCallback != null) {
            ackManager.registerAck(message.getMsgId(), getChannelId(), ackCallback);
        }
        Channel channel = channelHandlerContext.channel();
        channel.writeAndFlush(message);
    }

    @Override
    public <T extends Message> Future<T> sendCallback(M message, int timeout) {
        String msgId = message.getMsgId();
        ThreadUtil.execute(() -> send(message, new AckCallback<T>(timeout) {
            @Override
            public void onSuccess(T result) {
                callbackFlag.add(msgId);
                callbackReturn.put(msgId, result);
            }
            @Override
            public void onTimeout() {
                callbackFlag.add(msgId);
            }
        }));
        return ThreadUtil.execAsync(() -> {
            while (true) {
                if(callbackFlag.contains(msgId) && callbackReturn.containsKey(msgId)) {
                    callbackFlag.remove(msgId);
                    return (T) callbackReturn.get(callbackReturn);
                }
            }
        });
    }

    public <T extends Message> T sendAndReceive(M message, int timeout) {
        Future<T> future = sendCallback(message, timeout);
        try {
            future.get(timeout, TimeUnit.SECONDS);
            return (T) callbackReturn.remove(message.getMsgId());
        } catch (Exception e) {
            return null;
        }
    }

    public <T extends Message> T sendAndReceiveRetry(M message, int timeout, int retryNumber) {
        int count = 1;
        do {
            T result = sendAndReceive(message, timeout);
            if(result != null) {
                return result;
            } else {
                count++;
            }
        }while (count <= retryNumber);
        return null;
    }

    @Override
    public void send(M message) {
        send(message, null);
    }

    @Override
    public void broadcast(String clusterName, M message) {
        message.setDirection(Message.Direction.BROADCAST);
        pubSubStore.publish(clusterName, message);
    }

    @Override
    public void set(String key, Object value) {
        store.set(key, value);
    }

    @Override
    public <T> T get(String key) {
        return store.get(key);
    }

    @Override
    public boolean has(String key) {
        return store.has(key);
    }

    @Override
    public void del(String key) {
        store.del(key);
    }
}
